# 使用MNIST数据集对一个二层神经网络进行训练
# 并图形化显示随着学习，损失函数的变化

import sys,os
sys.path.append(os.pardir)
import numpy as np 
import matplotlib.pyplot as plt
from dataset.mnist import load_mnist
from two_layer_net import TwoLayerNet

# 读入数据
(x_train,t_train),(x_test,t_test) = load_mnist(normalize=True,one_hot_label=True)
# 初始化网络
network = TwoLayerNet(input_size=784,hidden_size=50,output_size=10)

# 初始化超参数
iters_num = 10000
train_size = x_train.shape[0]
batch_size = 100
learning_rate = 0.1

train_loss_list = []
train_acc_list = []
test_acc_list = []

iter_per_epoch = max(train_size/batch_size,1)

# 使用随机梯度下降算法进行训练
for i in range(iters_num):
    batch_mask = np.random.choice(train_size,batch_size)
    x_batch = x_train[batch_mask]
    t_batch = t_train[batch_mask]

    # 一个批量输入后，计算一次梯度
    grad = network.gradient(x_batch,t_batch)

    # 更新参数
    for key in ('W1','b1','W2','b2'):
        network.params[key] -= learning_rate*grad[key]

    # 计算损失
    loss = network.loss(x_batch,t_batch)
    train_loss_list.append(loss)

    if i % iter_per_epoch == 0:
        train_acc = network.accuracy(x_train,t_train)
        test_acc = network.accuracy(x_test,t_test)
        train_acc_list.append(train_acc)
        test_acc_list.append(test_acc)
        print("train acc, test acc |"+str(train_acc)+","+str(test_acc))

# 图形显示
# 1. 随着更新次数，损失函数的变化
# plt.plot(train_loss_list)
# plt.show()

# 2. 比较训练数据的精度和测试数据的精度
# markers = {'train':'o','test':'s'}
# x = np.arange(len(train_acc_list))
# plt.plot(x,train_acc_list,label='train acc')
# plt.plot(x,test_acc_list,label='test acc',linestyle='--')
# plt.xlabel("epochs")
# plt.ylabel("accuracy")
# plt.ylim(0,1.0)
# plt.legend(loc='lower right')

# plt.show()

# 使用训练得到的参数，对单个样本输入进行预测
y = network.predict(x_test[0])
print(np.argmax(y))
print(np.argmax(t_test[0]))
